import pickle
import numpy as np
import pandas as pd
import faiss
import joblib
import cv2
import gradio as gr
from sklearn.model_selection import train_test_split
from lazypredict.Supervised import LazyClassifier
from util import createXY

# 读取数据
X, y = createXY(train_folder="D:\\VScode\\vscode-python\\data\\train",
                dest_folder='', method='flat')
X = np.array(X).astype('float32')
faiss.normalize_L2(X)  # 对数据进行L2归一化
y = np.array(y)
# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.25, random_state=2023)
clf = LazyClassifier()
result, _ = clf.fit(X_train, X_test, y_train, y_test)
print(result)

# 保存最好的模型
best_model_name = result['F1 Score'].idxmax()  # 获取F1分数最高行的索引值，即：模型名称
print("\nF1分数最高的模型是: ", best_model_name)
best_model = clf.models[best_model_name]  # 根据模型名称，从模型字典中获取模型对象
result = best_model.predict(X_test)  # 该字典可以直接被拿来进行预测
print(f"用{best_model_name}预测X_test的结果是:\n{result}")
joblib.dump(best_model, open('best_model.pkl', 'wb'))

# 加载模型
model = joblib.load('best_model.pkl')


def predict_image(img):
    img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)  # 将图像从BGR格式转换为RGB格式
    img = cv2.resize(img, (32, 32))  # 调整图像大小到32x32
    img = img.reshape(1, -1)  # 将图像展平
    prediction = model.predict(img)
    return "Dog" if prediction[0] == 1 else "Cat"


# 创建 Gradio 界面
iface = gr.Interface(
    inputs="image",
    outputs="label",
    fn=predict_image
)

# 启动 Gradio 界面
iface.launch()  # 启动界面
